import requests
import base64
import os

# 先决决定的变量。
src_path = input('Source image path:\n>>> ').strip(
).removeprefix('"').removesuffix('"')
src_name, src_format = src_path.split(os.path.sep)[-1].split('.')
out_path = f'{os.path.dirname(src_path)}{os.path.sep}{src_name}-HD.{src_format}'

API_KEY = 'YDt7GDPQe2gLx6nYwtrUWGt4'
SECRET_KEY = 'OYImlDctSaUUP244cnEMcSF5v1Rt2BvU'

# 获取 token。
response = requests.get(
    f'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={API_KEY}&client_secret={SECRET_KEY}')

if response.status_code != 200:
    print(f'Request failed: {response.status_code}')
    exit(-1)

token = response.json()['access_token']

# 修复图片。
with open(src_path, 'rb') as fr:
    src_img = base64.b64encode(fr.read())

url = f'https://aip.baidubce.com/rest/2.0/image-process/v1/image_definition_enhance?access_token={token}'
headers = {
    'content-type': 'application/x-www-form-urlencoded',
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.159 Safari/537.36 Edg/92.0.902.84'
}
response = requests.post(url, headers=headers, data={"image": src_img})

if response.status_code != 200:
    print(f'Request failed: {response.status_code}')
    exit(-1)

out_img = base64.b64decode(response.json()['image'])
with open(out_path, 'wb') as fw:
    fw.write(out_img)
